import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
import net.resnet as res
from .gcn.gcn_model import gcn_model
import os
import time
import datetime
from .pcb_model_one import PCB

def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out')
        nn.init.constant_(m.bias, 0.0)
    elif classname.find('Conv') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
        if m.bias is not None:
            nn.init.constant_(m.bias, 0.0)
    elif classname.find('BatchNorm') != -1:
        if m.affine:
            nn.init.constant_(m.weight, 1.0)
            nn.init.constant_(m.bias, 0.0)

def weights_init_classifier(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.normal_(m.weight, std=0.001)
        if m.bias:
            nn.init.constant_(m.bias, 0.0)


#baseline model
class Resnet50_baseline(nn.Module):
    def __init__(self,pooling=True,stride=1):
        super(Resnet50_s1,self).__init__()
        original = models.resnet50(pretrained=True).state_dict()
        self.backbone = res.ResNet(last_stride=stride)
        for key in original:
            if key.find('fc') != -1:
                continue
            self.backbone.state_dict()[key].copy_(original[key])
        del original
      
        self.add_module('avgpool',nn.AdaptiveAvgPool2d(1))
        self.out_dim = 2048
    
    def forward(self,x):
        x = self.backbone(x)
        x = self.avgpool(x)
        x = x.view(x.shape[0],-1)
        return x

#our model
class Resnet50_stgcn(nn.Module):
    def __init__(self,pooling=True,stride=1,S=8,args=None):
        super(Resnet50_stgcn,self).__init__()
        original = models.resnet50(pretrained=True).state_dict()
        self.backbone = res.ResNet(last_stride=stride)
        self.S=args.S
        self.gcn_model=gcn_model(args)
        for key in original:
            if key.find('fc') != -1:
                continue
            self.backbone.state_dict()[key].copy_(original[key])
        del original
       
        self.add_module('avgpool',nn.AdaptiveAvgPool2d(1))

        self.out_dim = 2048
        self.tran=args.tran
        self.graph_arch=args.graph_arch
        self.fus_t=args.fus_t
    

    def forward(self,x):
        x = self.backbone(x)
        gcn_feat=self.gcn_model(x) # list
     
        ### if
        x = self.avgpool(x)
        x = x.view(x.shape[0],-1)
        x=torch.mean(x.reshape(x.shape[0]//self.S,self.S,-1),dim=1)

        gcn_feat.append(x)
        assert len(gcn_feat)==3 
        x=gcn_feat # list
        return x


class CNN(nn.Module):
    def __init__(self,out_dim,model_type='resnet50_s1',num_class=710,stride=1,args=None):
        super(CNN,self).__init__()
        self.model_type = model_type
        if model_type == 'resnet50_baseline':
            self.features = Resnet50_baseline(stride=stride)
        else:
            print("using stgcn")
            self.features = Resnet50_stgcn(stride=stride,S=args.S,args=args)
      
        
        if args.fusion_type!="concat":
            self.bottleneck = nn.BatchNorm1d(out_dim)
        else:
            self.bottleneck = nn.BatchNorm1d(out_dim+args.g_feat*args.gcn_num)
          
        self.bottleneck.bias.requires_grad_(False)  # no shift
        self.bottleneck.apply(weights_init_kaiming)

     
        if args.fusion_type!="concat":
            self.classifier = nn.Linear(out_dim,num_class, bias=False)
        else:
          
            self.classifier = nn.Linear(out_dim+args.g_feat*args.gcn_num,num_class, bias=False)
     
        self.classifier.apply(weights_init_classifier)
       
        self.dropout=args.dp
        self.bn_loc=args.bn_loc
        self.graph_arch=args.graph_arch
        self.loss_type=args.loss_type
        self.sole=args.sole
        self.fusion_type=args.fusion_type

    def forward(self,x,seg=None):
     
        x = self.features(x) # type: list  baseline:tensor 
        if(self.fusion_type=='mean'):
            pool_x=(x[0]+x[1]+x[2])/3
        else:
            x=x

        if(self.fusion_type!='concat'):
            bn = self.bottleneck(pool_x)
        else:
            if(self.loss_type=='split'):
                bn = self.bottleneck(torch.cat(x,dim=1))
            elif(self.loss_type=='concat'):
                x=torch.cat(x,dim=1)
                bn = self.bottleneck(x)
            else:
                bn = self.bottleneck(x) #baseline and newbaseline 

        if self.training == True:
            output = self.classifier(bn)
       
            if(self.bn_loc=='before'):
                return  x,output # x type: list or tensor.
            else:
                return bn,output
        else:
            if self.sole:
                return bn,orig_x
            else:
                return bn

if __name__ == '__main__':
    model = Resnet50_baseline(stride=stride)
    input = torch.ones(3,3,256,128)
    output = model(input)
    print(output.shape)
